from Ui_RoboticArm import Ui_MyMainWindow
from PyQt5.QtCore import QObject, pyqtSignal, QThread
from PyQt5.QtCore import QTimer
from PyQt5.QtWidgets import *
from PyQt5.QtSerialPort import QSerialPort, QSerialPortInfo
import sys
import os
from robotdriver import robot_driver, my_robot
import numpy as np
import time
import pyqtgraph.opengl as gl
from myThread import MyThread1, MyThread2
from blue_tooth import BluetoothDataTransfer

EPSINON = 1E-4

def get_link_points(tf_matrix):
    point_e1_ = np.array([0.0, 0.0, 20.0, 1.0])
    point_e2_ = np.array([0.0, 0.0, -20.0, 1.0])
    point_e1 = (np.dot(tf_matrix, point_e1_))[0:3]
    point_e2 = (np.dot(tf_matrix, point_e2_))[0:3]
    return np.array([point_e1, point_e2])

class MainWindow(QMainWindow, Ui_MyMainWindow):
    def __init__(self, parent = None):
        super(MainWindow, self).__init__(parent)
        self.setupUi(self)
        self.setFixedSize(1200, 900) 
        #   init sim_mode
        self.radioButton_sim_mode.setChecked(True)
        self.createItems()
        self.createSignalSlot()
        self.set_graph_widget()
        self.update_text_thread_sim = None
        self.update_text_thread_real = None
        self.bluetoothTransfer = None
        self.Bluetooth_device_dict = {}
        self.update_mode()
        
    def terminate_thread_sim(self):
        if(self.update_text_thread_sim !=None):
            self.update_text_thread_sim.stop()
            self.update_text_thread_sim.terminate()
            self.update_text_thread_sim = None
            
    def create_thread_sim(self):
        if(self.update_text_thread_sim != None):
            self.terminate_thread_sim()
        self.update_text_thread_sim = MyThread1()
        self.update_text_thread_sim.update_data.connect(self.update_text_area)
        self.update_text_thread_sim.update_data.connect(self.update_graph_widget)
        self.update_text_thread_sim.start()
        
    def terminate_thread_real(self):
        if(self.update_text_thread_real !=None):
            self.update_text_thread_real.stop()
            self.update_text_thread_real.terminate()
            self.update_text_thread_real = None

    def create_thread_real(self, bd):
        if(self.update_text_thread_real != None):
            self.terminate_thread_real()
        self.update_text_thread_real = MyThread2(bd)
        self.update_text_thread_real.update_data.connect(self.update_text_area)
        self.update_text_thread_real.update_data.connect(self.update_graph_widget)
        self.update_text_thread_real.start()
            
    def update_mode(self):
        if(self.radioButton_sim_mode.isChecked()):
            self.pushButtonJ1P.setEnabled(True)
            self.pushButtonJ1N.setEnabled(True)
            self.pushButtonJ2P.setEnabled(True)
            self.pushButtonJ2N.setEnabled(True)
            self.pushButtonJ3P.setEnabled(True)
            self.pushButtonJ3N.setEnabled(True)
            self.pushButtonJ4P.setEnabled(True)
            self.pushButtonJ4N.setEnabled(True)
            self.pushButtonJ5P.setEnabled(True)
            self.pushButtonJ5N.setEnabled(True)
            self.pushButtonJ6P.setEnabled(True)
            self.pushButtonJ6N.setEnabled(True)
            self.pushButton_XP.setEnabled(True)
            self.pushButton_XN.setEnabled(True)
            self.pushButton_YP.setEnabled(True)
            self.pushButton_YN.setEnabled(True)
            self.pushButton_ZP.setEnabled(True)
            self.pushButton_ZN.setEnabled(True)
            self.pushButtonHome.setEnabled(True)
            self.pushButton_connect_device.setEnabled(False)
            self.pushButton_disconnect_device.setEnabled(False)
            self.pushButton_rescan.setEnabled(False)
            self.terminate_thread_real()
            self.create_thread_sim()
        elif (self.radioButton_real_mode.isChecked()):
            self.pushButtonJ1P.setEnabled(False)
            self.pushButtonJ1N.setEnabled(False)
            self.pushButtonJ2P.setEnabled(False)
            self.pushButtonJ2N.setEnabled(False)
            self.pushButtonJ3P.setEnabled(False)
            self.pushButtonJ3N.setEnabled(False)
            self.pushButtonJ4P.setEnabled(False)
            self.pushButtonJ4N.setEnabled(False)
            self.pushButtonJ5P.setEnabled(False)
            self.pushButtonJ5N.setEnabled(False)
            self.pushButtonJ6P.setEnabled(False)
            self.pushButtonJ6N.setEnabled(False)
            self.pushButton_XP.setEnabled(False)
            self.pushButton_XN.setEnabled(False)
            self.pushButton_YP.setEnabled(False)
            self.pushButton_YN.setEnabled(False)
            self.pushButton_ZP.setEnabled(False)
            self.pushButton_ZN.setEnabled(False)
            self.pushButtonHome.setEnabled(False)
            self.pushButton_connect_device.setEnabled(True)
            self.pushButton_disconnect_device.setEnabled(True)
            self.pushButton_rescan.setEnabled(True)
            self.terminate_thread_sim()
        
    def set_graph_widget(self):
        self.graphView.show()
        self.graphView.setWindowTitle('MY 6-DOF RoboticArm')
        self.graphView.setCameraPosition(distance=1200, azimuth=-45)
        # Grid
        gx = gl.GLGridItem()
        gx.setSize(1000, 1000, 1)
        gx.setSpacing(50, 50, 1)
        self.graphView.addItem(gx)
        # Axis
        axisitem = gl.GLAxisItem()
        axisitem.setSize(x=300, y=300, z=300)
        self.graphView.addItem(axisitem)
        # Text X
        txt_x = gl.GLTextItem()
        txt_x.setData(pos=(300, 0, 0), color=(127, 255, 127, 255), text='X')
        self.graphView.addItem(txt_x)
        # Text Y
        txt_y = gl.GLTextItem()
        txt_y.setData(pos=(0, 300, 0), color=(127, 255, 127, 255), text='Y')
        self.graphView.addItem(txt_y)
        # Text Y
        txt_z = gl.GLTextItem()
        txt_z.setData(pos=(0, 0, 300), color=(127, 255, 127, 255), text='Z')
        self.graphView.addItem(txt_z)
        #####################################################################
        joints = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
        robot_tf_tree = my_robot.get_robot_tf_tree(joints)
        ### base_box
        tf_matrix_base_0 = my_robot.get_base_robot_tf()
        base_origin = (0, 0, 0)
        quat_base = (0, 0, 0, 1)
        ### link 1
        tf_matrix_base_1 = np.dot(robot_tf_tree[0], tf_matrix_base_0)
        link1_origin = (tf_matrix_base_1[0,3], tf_matrix_base_1[1,3], tf_matrix_base_1[2,3])
        ### link 2
        tf_matrix_base_2 = np.dot(tf_matrix_base_1, robot_tf_tree[1])
        link2_origin = (tf_matrix_base_2[0,3], tf_matrix_base_2[1,3], tf_matrix_base_2[2,3])
        ### link 3
        tf_matrix_base_3 = np.dot(tf_matrix_base_2, robot_tf_tree[2])
        link3_origin = (tf_matrix_base_3[0,3], tf_matrix_base_3[1,3], tf_matrix_base_3[2,3])
        ### link 4
        tf_matrix_base_4 = np.dot(tf_matrix_base_3, robot_tf_tree[3])
        link4_origin = (tf_matrix_base_4[0,3], tf_matrix_base_4[1,3], tf_matrix_base_4[2,3])
        ### link 5
        tf_matrix_base_5 = np.dot(tf_matrix_base_4, robot_tf_tree[4])
        link5_origin = (tf_matrix_base_5[0,3], tf_matrix_base_5[1,3], tf_matrix_base_5[2,3])
        ### link 6
        tf_matrix_base_6 = np.dot(tf_matrix_base_5, robot_tf_tree[5])
        link6_origin = (tf_matrix_base_6[0,3], tf_matrix_base_6[1,3], tf_matrix_base_6[2,3])
        ### TCP
        tf_matrix_base_tcp = np.dot(tf_matrix_base_6, my_robot.get_end_tcp_tf())
        tcp_origin = (tf_matrix_base_tcp[0,3], tf_matrix_base_tcp[1,3], tf_matrix_base_tcp[2,3])
        #####################################
        pos = np.array([base_origin, 
                    link1_origin, 
                    link2_origin, 
                    link3_origin, 
                    link4_origin, 
                    link5_origin, 
                    link6_origin, 
                    tcp_origin])
        size = np.array([10, 10, 10, 10, 10, 10, 10, 10])
        color = np.array([(1.0, 1.0, 0.0, 0.5),
                            (1.0, 1.0, 0.0, 0.5),
                            (1.0, 1.0, 0.0, 0.5),
                            (1.0, 1.0, 0.0, 0.5),
                            (1.0, 1.0, 0.0, 0.5),
                            (1.0, 1.0, 0.0, 0.5),
                            (1.0, 1.0, 0.0, 0.5),
                            (1.0, 1.0, 0.0, 0.5)])
        self.sp_point = gl.GLScatterPlotItem(pos=pos, size=size, color=color, pxMode=False)
        self.graphView.addItem(self.sp_point)
        color = np.array([(1.0, 1.0, 1.0, 1),
                            (1.0, 1.0, 1.0, 1),
                            (1.0, 1.0, 1.0, 1),
                            (1.0, 1.0, 1.0, 1),
                            (1.0, 1.0, 1.0, 1),
                            (1.0, 1.0, 1.0, 1),
                            (1.0, 1.0, 1.0, 1),
                            (1.0, 1.0, 1.0, 1)])
        self.lp_line = gl.GLLinePlotItem(pos=pos, color=color, width=5, antialias=True)
        self.graphView.addItem(self.lp_line)
        ##
        color = np.array([(0, 191/255.0, 1, 1), (0, 191/255.0, 1, 1)])
        pos_link1 = get_link_points(tf_matrix_base_1)
        self.lp_link1 = gl.GLLinePlotItem(pos=pos_link1, color=color, width=80, antialias=False)
        pos_link2 = get_link_points(tf_matrix_base_2)
        self.lp_link2 = gl.GLLinePlotItem(pos=pos_link2, color=color, width=20, antialias=True)
        pos_link3 = get_link_points(tf_matrix_base_3)
        self.lp_link3 = gl.GLLinePlotItem(pos=pos_link3, color=color, width=20, antialias=True)
        pos_link4 = get_link_points(tf_matrix_base_4)
        self.lp_link4 = gl.GLLinePlotItem(pos=pos_link4, color=color, width=20, antialias=True)
        pos_link5 = get_link_points(tf_matrix_base_5)
        self.lp_link5 = gl.GLLinePlotItem(pos=pos_link5, color=color, width=20, antialias=True)
        pos_link6 = get_link_points(tf_matrix_base_6)
        self.lp_link6 = gl.GLLinePlotItem(pos=pos_link6, color=color, width=20, antialias=True)
        self.graphView.addItem(self.lp_link1)
        self.graphView.addItem(self.lp_link2)
        self.graphView.addItem(self.lp_link3)
        self.graphView.addItem(self.lp_link4)
        self.graphView.addItem(self.lp_link5)
        self.graphView.addItem(self.lp_link6)
        ##########
        
    def update_graph_widget(self, joints):
        if joints == None:
            return
        robot_tf_tree = my_robot.get_robot_tf_tree(joints)
        ### base_box
        tf_matrix_base_0 = my_robot.get_base_robot_tf()
        base_origin = (0, 0, 0)
        quat_base = (0, 0, 0, 1)
        ### link 1
        tf_matrix_base_1 = np.dot(robot_tf_tree[0], tf_matrix_base_0)
        link1_origin = (tf_matrix_base_1[0,3], tf_matrix_base_1[1,3], tf_matrix_base_1[2,3])
        ### link 2
        tf_matrix_base_2 = np.dot(tf_matrix_base_1, robot_tf_tree[1])
        link2_origin = (tf_matrix_base_2[0,3], tf_matrix_base_2[1,3], tf_matrix_base_2[2,3])
        ### link 3
        tf_matrix_base_3 = np.dot(tf_matrix_base_2, robot_tf_tree[2])
        link3_origin = (tf_matrix_base_3[0,3], tf_matrix_base_3[1,3], tf_matrix_base_3[2,3])
        ### link 4
        tf_matrix_base_4 = np.dot(tf_matrix_base_3, robot_tf_tree[3])
        link4_origin = (tf_matrix_base_4[0,3], tf_matrix_base_4[1,3], tf_matrix_base_4[2,3])
        ### link 5
        tf_matrix_base_5 = np.dot(tf_matrix_base_4, robot_tf_tree[4])
        link5_origin = (tf_matrix_base_5[0,3], tf_matrix_base_5[1,3], tf_matrix_base_5[2,3])
        ### link 6
        tf_matrix_base_6 = np.dot(tf_matrix_base_5, robot_tf_tree[5])
        link6_origin = (tf_matrix_base_6[0,3], tf_matrix_base_6[1,3], tf_matrix_base_6[2,3])
        ### TCP
        tf_matrix_base_tcp = np.dot(tf_matrix_base_6, my_robot.get_end_tcp_tf())
        tcp_origin = (tf_matrix_base_tcp[0,3], tf_matrix_base_tcp[1,3], tf_matrix_base_tcp[2,3])
        #####################################
        pos = np.array([base_origin, 
                    link1_origin, 
                    link2_origin, 
                    link3_origin, 
                    link4_origin, 
                    link5_origin, 
                    link6_origin, 
                    tcp_origin])
        pos_link1 = get_link_points(tf_matrix_base_1)
        pos_link2 = get_link_points(tf_matrix_base_2)
        pos_link3 = get_link_points(tf_matrix_base_3)
        pos_link4 = get_link_points(tf_matrix_base_4)
        pos_link5 = get_link_points(tf_matrix_base_5)
        pos_link6 = get_link_points(tf_matrix_base_6)
        
        self.sp_point.setData(pos=pos)
        self.lp_line.setData(pos=pos)
        self.lp_link1.setData(pos=pos_link1)
        self.lp_link2.setData(pos=pos_link2)
        self.lp_link3.setData(pos=pos_link3)
        self.lp_link4.setData(pos=pos_link4)
        self.lp_link5.setData(pos=pos_link5)
        self.lp_link6.setData(pos=pos_link6)
        
    def update_text_area(self, joints):
        if robot_driver.is_ready() or self.radioButton_real_mode.isChecked():
            current_joints = joints[:]
            current_pose = robot_driver.kinesolver.forward_kinematics(current_joints)
            current_euler = current_pose.pose_to_euler()
            for i in range(len(current_joints)):
                current_joints[i] = round(current_joints[i], 5) if abs(current_joints[i]) > EPSINON else 0.00
            for i in range(len(current_euler)):
                current_euler[i] = round(current_euler[i], 5) if abs(current_euler[i]) > EPSINON else 0.00
            current_pose.x = round(current_pose.x, 5) if abs(current_pose.x) > EPSINON else 0.00
            current_pose.y = round(current_pose.y, 5) if abs(current_pose.y) > EPSINON else 0.00
            current_pose.z = round(current_pose.z, 5) if abs(current_pose.z) > EPSINON else 0.00
            current_pose.qx = round(current_pose.qx, 5) if abs(current_pose.qx) > EPSINON else 0.00
            current_pose.qy = round(current_pose.qy, 5) if abs(current_pose.qy) > EPSINON else 0.00
            current_pose.qz = round(current_pose.qz, 5) if abs(current_pose.qz) > EPSINON else 0.00
            current_pose.qw = round(current_pose.qw, 5) if abs(current_pose.qw) > EPSINON else 0.00
            
            if self.radioButton_degree.isChecked():
                self.textEdit_j1.setText(str(np.rad2deg(current_joints[0])))
                self.textEdit_j2.setText(str(np.rad2deg(current_joints[1])))
                self.textEdit_j3.setText(str(np.rad2deg(current_joints[2])))
                self.textEdit_j4.setText(str(np.rad2deg(current_joints[3])))
                self.textEdit_j5.setText(str(np.rad2deg(current_joints[4])))
                self.textEdit_j6.setText(str(np.rad2deg(current_joints[5])))
            elif self.radioButton_rad.isChecked():
                self.textEdit_j1.setText(str(current_joints[0]))
                self.textEdit_j2.setText(str(current_joints[1]))
                self.textEdit_j3.setText(str(current_joints[2]))
                self.textEdit_j4.setText(str(current_joints[3]))
                self.textEdit_j5.setText(str(current_joints[4]))
                self.textEdit_j6.setText(str(current_joints[5]))
            
            self.textEdit_X.setText(str(current_pose.x))
            self.textEdit_Y.setText(str(current_pose.y))
            self.textEdit_Z.setText(str(current_pose.z))
            self.textEdit_qx.setText(str(current_pose.qx))
            self.textEdit_qy.setText(str(current_pose.qy))
            self.textEdit_qz.setText(str(current_pose.qz))
            self.textEdit_qw.setText(str(current_pose.qw))
            
            self.textEdit_rx.setText(str(current_euler[0]))
            self.textEdit_ry.setText(str(current_euler[1]))
            self.textEdit_rz.setText(str(current_euler[2]))
        
    def createItems(self):
        pass

    def createSignalSlot(self):
        self.pushButtonJ1P.clicked.connect(self.set_joint1_movej_p)
        self.pushButtonJ1N.clicked.connect(self.set_joint1_movej_n)
        self.pushButtonJ2P.clicked.connect(self.set_joint2_movej_p)
        self.pushButtonJ2N.clicked.connect(self.set_joint2_movej_n)  
        self.pushButtonJ3P.clicked.connect(self.set_joint3_movej_p)
        self.pushButtonJ3N.clicked.connect(self.set_joint3_movej_n)  
        self.pushButtonJ4P.clicked.connect(self.set_joint4_movej_p)
        self.pushButtonJ4N.clicked.connect(self.set_joint4_movej_n)  
        self.pushButtonJ5P.clicked.connect(self.set_joint5_movej_p)
        self.pushButtonJ5N.clicked.connect(self.set_joint5_movej_n)  
        self.pushButtonJ6P.clicked.connect(self.set_joint6_movej_p)
        self.pushButtonJ6N.clicked.connect(self.set_joint6_movej_n)
        self.pushButton_XP.clicked.connect(self.set_cartpose_x_movej_p)
        self.pushButton_XN.clicked.connect(self.set_cartpose_x_movej_n)
        self.pushButton_YP.clicked.connect(self.set_cartpose_y_movej_p)
        self.pushButton_YN.clicked.connect(self.set_cartpose_y_movej_n)
        self.pushButton_ZP.clicked.connect(self.set_cartpose_z_movej_p)
        self.pushButton_ZN.clicked.connect(self.set_cartpose_z_movej_n)
        self.pushButtonHome.clicked.connect(self.set_robot_home)
        
        self.radioButton_real_mode.clicked.connect(self.update_mode)
        self.radioButton_sim_mode.clicked.connect(self.update_mode)
        
        self.pushButton_rescan.clicked.connect(self.scan_bt_device)
        self.pushButton_connect_device.clicked.connect(self.connect_bt_device)
        
    def scan_bt_device(self):
        devices = BluetoothDataTransfer.scan_devices()
        self.Bluetooth_device_dict = {}
        info = []
        for addr, name in devices:
            info.append(name)
            self.Bluetooth_device_dict[name] = addr
        self.comboBox_bt_devices.addItems(info)
        
    def connect_bt_device(self):
        if self.bluetoothTransfer != None:
            try:
                self.bluetoothTransfer.disconnect()
            except:
                pass
        if self.Bluetooth_device_dict == {}:
            print("111")
            return
        device_name = self.comboBox_bt_devices.currentText()
        if device_name not in self.Bluetooth_device_dict.keys():
            print("222")
            return
        device_addr = self.Bluetooth_device_dict[device_name]
        self.bluetoothTransfer = BluetoothDataTransfer(device_addr, device_name, 1)
        if self.bluetoothTransfer.connect():
            self.checkBox_connect_status.setChecked(True)
        else:
            self.checkBox_connect_status.setChecked(False)
            print("333")
            return
        self.create_thread_real(self.bluetoothTransfer)
        
    def disconnect_bt_device(self):
        if self.bluetoothTransfer == None:
            return
        self.bluetoothTransfer = None
        self.terminate_thread_real()
        self.checkBox_connect_status.setChecked(False)
        
    def set_joint1_movej_p(self):
        np.deg2rad(self.SpinBoxSetJ1.value())
        robot_driver.move_joint1_slightly(np.deg2rad(self.SpinBoxSetJ1.value()))
    def set_joint2_movej_p(self):
        robot_driver.move_joint2_slightly(np.deg2rad(self.SpinBoxSetJ2.value()))
    def set_joint3_movej_p(self):
        robot_driver.move_joint3_slightly(np.deg2rad(self.SpinBoxSetJ3.value()))
    def set_joint4_movej_p(self):
        robot_driver.move_joint4_slightly(np.deg2rad(self.SpinBoxSetJ4.value()))
    def set_joint5_movej_p(self):
        robot_driver.move_joint5_slightly(np.deg2rad(self.SpinBoxSetJ5.value()))
    def set_joint6_movej_p(self):
        robot_driver.move_joint6_slightly(np.deg2rad(self.SpinBoxSetJ6.value()))
    def set_cartpose_x_movej_p(self):
        robot_driver.move_x_slightly(self.SpinBox_pos_x.value())
    def set_cartpose_y_movej_p(self):
        robot_driver.move_y_slightly(self.SpinBox_pos_y.value())
    def set_cartpose_z_movej_p(self):
        robot_driver.move_z_slightly(self.SpinBox_pos_z.value())
        
    def set_joint1_movej_n(self):
        robot_driver.move_joint1_slightly(np.deg2rad(-1.0*self.SpinBoxSetJ1.value()))
    def set_joint2_movej_n(self):
        robot_driver.move_joint2_slightly(np.deg2rad(-1.0*self.SpinBoxSetJ2.value()))
    def set_joint3_movej_n(self):
        robot_driver.move_joint3_slightly(np.deg2rad(-1.0*self.SpinBoxSetJ3.value()))
    def set_joint4_movej_n(self):
        robot_driver.move_joint4_slightly(np.deg2rad(-1.0*self.SpinBoxSetJ4.value()))
    def set_joint5_movej_n(self):
        robot_driver.move_joint5_slightly(np.deg2rad(-1.0*self.SpinBoxSetJ5.value()))
    def set_joint6_movej_n(self):
        robot_driver.move_joint6_slightly(np.deg2rad(-1.0*self.SpinBoxSetJ6.value()))
    def set_cartpose_x_movej_n(self):
        robot_driver.move_x_slightly(-1.0*self.SpinBox_pos_x.value())
    def set_cartpose_y_movej_n(self):
        robot_driver.move_y_slightly(-1.0*self.SpinBox_pos_y.value())
    def set_cartpose_z_movej_n(self):
        robot_driver.move_z_slightly(-1.0*self.SpinBox_pos_z.value())

    def set_robot_home(self):
        theta_1 = 0.0
        theta_2 = -np.pi/2
        theta_3 = 0.0
        theta_4 = 0.0
        theta_5 = 0.0
        theta_6 = 0.0
        joints = [theta_1, theta_2, theta_3, theta_4, theta_5, theta_6]
        robot_driver.move_to_home(joints)

# window = QApplication(sys.argv)
# my_robot_win = MainWindow()
# # my_robot_win.show()
# sys.exit(window.exec_())

if __name__ == '__main__':
    window = QApplication(sys.argv)
    TheWin = MainWindow()
    TheWin.show()
    sys.exit(window.exec_())